from ultralytics import YOLO
import yaml
import os
import shutil

def train_model():
    # 加载预训练模型
    model = YOLO("yolo11n.pt")  # 使用YOLO模型作为基础

    # 训练配置
    config = {
        "data": "data.yaml",
        "epochs": 150,  # 增加训练轮次
        "patience": 40,  # 延长早停等待
        "batch": 32 if ... else 16,  # 减小batch size提高稳定性
        "lr0": 0.0005,  # 初始学习率再降低50%
        "lrf": 0.005,   # 最终学习率减半
        "optimizer": "SGD",  # 换回SGD+momentum(更适合目标检测)
        "momentum": 0.937,  # 添加动量参数
        "warmup_epochs": 10,  # 延长热身期
        # 数据增强优化
        "mixup": 0.0,    # 关闭mixup(车辆场景不适用)
        "copy_paste": 0.0, # 关闭copy-paste(会产生不合理车辆组合)
        "degrees": 5.0,   # 减小旋转角度(避免颠倒的车辆)
        "shear": 1.0,     # 减小剪切幅度
        "erasing": 0.2,   # 降低随机擦除概率
        # 新增关键参数
        "label_smoothing": 0.1,  # 减少过拟合
        "box": 7.5,       # 增加box loss权重(提高定位能力)
        "cls": 0.8,       # 降低cls loss权重(缓解置信度问题)
        "close_mosaic": 20,  # 提前关闭mosaic增强
        "weight_decay": 0.001  # 减小正则化强度
    }

    # 开始训练
    results = model.train(**config)

    # 自动将最佳模型复制为model.pt
    best_pt_path = f"runs/detect/{config['name']}/weights/best.pt"
    if os.path.exists(best_pt_path):
        shutil.copy(best_pt_path, "model.pt")
        print(f"✅ 最佳模型已保存为: model.pt")
    else:
        print(f"❌ 错误：未找到训练输出文件 {best_pt_path}")


if __name__ == "__main__":
    import torch
    train_model()